import sys 
sys.path.append('/home/aistudio/external-libraries')

import deepxde as dde
import paddle
import numpy as np
from matplotlib import pyplot as plt
import os

main_path = os.getcwd()

# 需要对哪个level的网格进行加密？
level = 1

"""------------------------------ Fluid parameters  -----------------------"""# 用于无量纲化的速度、压力和密度，取入口处的声速值、压力值和密度值
rho = 1.2 #(kg/m^3)
mu = 1.8e-5 #(Pa·s)
nu = mu/rho #(m^2/s)

"""------------------------------------ Training Spatio-temporal region --------------------------------------"""
"""-------------- Rectangular -------------"""
x_min, x_max = 0, 20
L = x_max - x_min

y_min, y_max = 0, 100
W = y_max - y_min

z_min, z_max = 0, 20
H = z_max - z_min

"""Spatial region definition"""
geom_box = dde.geometry.geometry_3d.Cuboid([x_min, y_min, z_min], [x_max, y_max, z_max])  # 设置左下角和右上角坐标
geom_truck = dde.geometry.geometry_3d.Cuboid([0, 80, 18.5], [4, 90, 20])
geom = dde.geometry.CSGDifference(geom_box, geom_truck)

'''----------------------- Other functions ---------------------------'''
def plot_points_3d(points):
    fig = plt.figure(dpi=300)
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0.5)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    # plt.axis("scaled") 
    plt.savefig('./Figs/Total_points.png', dpi=300)
    plt.show()

def reorder_points(coordinates):
    # 计算中心
    center_x = coordinates[:, 0].mean()
    center_y = coordinates[:, 1].mean()
    # 计算角度
    angle = np.arctan2(coordinates[:, 1] - center_y, coordinates[:, 0] - center_x)
    # 按角度值排序
    points_indices = np.argsort(angle)
    ordered_points = coordinates[points_indices, :]
    return ordered_points

def plot_flow_field():
    import matplotlib as mpl
    from matplotlib import rcParams
    config = {"font.size": 10}
    rcParams.update(config)
    plt.rcParams['xtick.direction'] = 'in'
    plt.rcParams['ytick.direction'] = 'in'
    
    '''------------------------------------------- Model validation ---------------------------------------------'''
    '''------------------Plotting 2------------------'''
    interface = dde.geometry.Rectangle([x_min, y_min], [x_max, y_max])
    interface_xy = interface.uniform_points(20000)

    test_data = np.hstack((interface_xy, 20 * np.ones_like(interface_xy[:, 0:1])))
    test_data2 = data.transform_inputs(test_data)

    # Model predictions generation
    u = model.predict(test_data2)[:, 0]
    v = model.predict(test_data2)[:, 1]
    w = model.predict(test_data2)[:, 2]
    p = model.predict(test_data2)[:, 3]
    velocity = (u**2 + v**2 + w**2)**0.5

    cmap_color = "jet"
    y_sticks = [-100, -80, -60, -40, -20, 0]
    y_labels = [100, 80, 60, 40, 20, 0]
    truck_plot = reorder_points(dde.geometry.Rectangle([-90, 0], [-80, 4]).uniform_boundary_points(500))

    # plot
    fig, ax = plt.subplots(figsize=(6, 1.5), dpi=300)
    ax.tricontourf(-test_data[:, 1], test_data[:, 0], velocity, levels=1000, cmap=cmap_color, vmin=0, vmax=9)
    ax.set_xticks(y_sticks)
    ax.set_xticklabels(y_labels)
    plt.axis('scaled')
    plt.fill(truck_plot[:, 0], truck_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=0, vmax=9)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('y')
    plt.ylabel('x')
    plt.title('Velocity Magnitude')
    plt.savefig("./Figs/Velocity.png")
    plt.show()   
    
    fig, ax = plt.subplots(figsize=(6, 1.5), dpi=300)
    ax.tricontourf(-test_data[:, 1], test_data[:, 0], p, levels=1000, cmap=cmap_color, vmin=-20, vmax=30)
    ax.set_xticks(y_sticks)
    ax.set_xticklabels(y_labels)
    plt.axis('scaled')
    plt.fill(truck_plot[:, 0], truck_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=-20, vmax=30)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[Pa]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('y')
    plt.ylabel('x')
    plt.title('Pressure')
    plt.savefig("./Figs/Pressure.png")
    plt.show()  
    
    fig, ax = plt.subplots(figsize=(6, 1.5), dpi=300)
    ax.tricontourf(-test_data[:, 1], test_data[:, 0], u, levels=1000, cmap=cmap_color, vmin=-0.705, vmax=1.15)
    ax.set_xticks(y_sticks)
    ax.set_xticklabels(y_sticks)
    plt.axis('scaled')
    plt.fill(truck_plot[:, 0], truck_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=-0.705, vmax=1.15)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('y')
    plt.ylabel('x')
    plt.title('u')
    plt.savefig("./Figs/u.png")
    plt.show()  

    fig, ax = plt.subplots(figsize=(6, 1.5), dpi=300)
    ax.tricontourf(-test_data[:, 1], test_data[:, 0], v, levels=1000, cmap=cmap_color, vmin=-9.13, vmax=0)
    ax.set_xticks(y_sticks)
    ax.set_xticklabels(y_sticks)
    plt.axis('scaled')
    plt.fill(truck_plot[:, 0], truck_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=-9.13, vmax=0)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('y')
    plt.ylabel('x')
    plt.title('v')
    plt.savefig("./Figs/v.png")
    plt.show()  

    fig, ax = plt.subplots(figsize=(6, 1.5), dpi=300)
    ax.tricontourf(-test_data[:, 1], test_data[:, 0], w, levels=1000, cmap=cmap_color, vmin=-0.5, vmax=0.5)
    ax.set_xticks(y_sticks)
    ax.set_xticklabels(y_sticks)
    plt.axis('scaled')
    plt.fill(truck_plot[:, 0], truck_plot[:, 1], color = 'w')
    norm =mpl.colors.Normalize(vmin=-0.5, vmax=0.5)
    plt.rcParams['ytick.direction'] = 'out'
    plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap_color), ax=ax, label='[m/s]')
    plt.rcParams['ytick.direction'] = 'in'    
    plt.xlabel('y')
    plt.ylabel('x')
    plt.title('w')
    plt.savefig("./Figs/w.png")
    plt.show()   

def make_dir():
    # 保存训练结果需要的文件
    # 模型保存地址
    Model_path = "Model/"  
    # 图片保存地址
    Figs_path = "Figs/"  
    # Loss保存地址
    Train_history_path = "Train_history/"  

    if not os.path.exists(Model_path):
    	os.mkdir(Model_path)
    if not os.path.exists(Figs_path):
    	os.mkdir(Figs_path)
    if not os.path.exists(Train_history_path):
    	os.mkdir(Train_history_path)

def plot_loss_history(loss_history, fname=None):
    loss_train = np.sum(loss_history.loss_train, axis=1)
    loss_test = np.sum(loss_history.loss_test, axis=1)

    plt.figure()
    plt.semilogy(loss_history.steps, loss_train, label="Train loss")
    plt.semilogy(loss_history.steps, loss_test, label="Test loss")
    for i in range(len(loss_history.metrics_test[0])):
        plt.semilogy(
            loss_history.steps,
            np.array(loss_history.metrics_test)[:, i],
            label="Test metric",
        )
    plt.xlabel("# Steps")
    plt.legend()

    if isinstance(fname, str):
        plt.savefig(fname, dpi=300)


'''-------------------------------------------------------------- training model:  ------------------------------------------------------------------'''
dde.config.set_random_seed(1234)
dde.config.set_default_float("float32")

make_dir()

'''------------------BC------------------'''
# data
data_points = np.load('../1_Fluent_results/level' + str(level) + '_fluent_results.npy')
data_xyz = data_points[:, 0:3]

'''--------------- Training Model ----------------'''
plot_points_3d(data_xyz)

'''----Training datasets and Loss function ----'''
data = dde.data.DataSet(X_train=data_xyz,
    					y_train=data_points[:, 3:7],
			    		X_test=data_xyz,
    					y_test=data_points[:, 3:7],
    					standardize=True,
						)  
print(data_points.shape)
'''--------- Neural Network setup ---------'''
layer_size = [3] + [40] * 8 + [4]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)

'''-------- Compile, Training and save Model --------'''
model = dde.Model(data, net)
weights = [4] * 1

# first use Adam to find a roughly good solution, 
model.compile("adam", lr=1e-3, loss_weights = weights)
loss_history, train_state = model.train(epochs=20000, display_every=1000)
# plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
# model.save(save_path="./Model/model")
plot_flow_field()

model.compile("adam", lr=5e-4, loss_weights = weights)
loss_history, train_state = model.train(epochs=20000, display_every=1000)
# plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
# model.save(save_path="./Model/model")
plot_flow_field()

model.compile("adam", lr=1e-4, loss_weights = weights)
loss_history, train_state = model.train(epochs=20000, display_every=1000)
# plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
# model.save(save_path="./Model/model")
plot_flow_field()

# and then use L-BFGS to search a local minimum.
dde.optimizers.config.set_LBFGS_options(maxcor=100, ftol=0, gtol=1e-08, maxiter=15000, maxfun=None, maxls=50)
model.compile(optimizer="L-BFGS", loss_weights = weights)
# Optional: Save the model during training.
checkpointer = dde.callbacks.ModelCheckpoint("Model/model", verbose=0, save_better_only=True)
loss_history, train_state = model.train(display_every=100, 
                                        # callbacks=[checkpointer],
                                        )
# plot_loss_history(loss_history, "./Train_history/Loss.png")
dde.saveplot(loss_history, train_state, issave=True, isplot=True, output_dir = "./Train_history")
# model.save(save_path="./Model/model")
plot_flow_field()
